import { z } from "zod";
import { CodeIndexOrchestrator } from "../services/orchestrator.js";
import { Translator } from "../services/translator.js";
import { Reranker } from "../services/reranker.js";
import path from "path";

export const searchCodeSchema = z.object({
  query: z.string().describe("Natural language or keyword search query"),
  directory_prefix: z.string().optional().describe("Optional directory path to filter results (e.g., 'src/components')"),
  max_results: z.number().int().positive().default(10).describe("Maximum number of results to return per category (default: 10)"),
  min_score: z.number().min(0).max(1).default(0.7).describe("Minimum similarity score threshold (0-1, default: 0.7)"),
  categorized_search: z.boolean().default(true).describe("Enable categorized search with separate results for text and code files (default: true)"),
  text_min_score: z.number().min(0).max(1).optional().describe("Minimum score for text files (default: 0.6)"),
  code_min_score: z.number().min(0).max(1).optional().describe("Minimum score for code files (default: 0.3)"),
  rerank: z.boolean().default(false).describe("Enable reranking of search results for better relevance (default: false)"),
});

export type SearchCodeArgs = z.infer<typeof searchCodeSchema>;

// Function to detect if a query is likely Chinese
function isChinese(query: string): boolean {
  return /["\u4e00"-"\u9fa5"]/.test(query);
}

export async function searchCodeTool(
  args: SearchCodeArgs,
  orchestrator: CodeIndexOrchestrator
): Promise<string> {
  try {
    const translator = new Translator();
    const translationMemoryPath = path.join(process.cwd(), 'translation-memory.json');
    await translator.loadTranslationMemory(translationMemoryPath);

    const originalQuery = args.query;
    const queryIsChinese = isChinese(originalQuery);
    const targetLang = queryIsChinese ? "English" : "Chinese";
    const translatedQuery = await translator.translate(originalQuery, targetLang);

    console.log(`Original query: "${originalQuery}"`);
    if (originalQuery !== translatedQuery) {
      console.log(`Translated query: "${translatedQuery}"`);
    }

    const searchTasks = [
      orchestrator.searchIndexWithCategories(
        originalQuery,
        args.directory_prefix,
        args.max_results,
        args.min_score,
        args.text_min_score,
        args.code_min_score
      ),
      orchestrator.searchIndexWithCategories(
        translatedQuery,
        args.directory_prefix,
        args.max_results,
        args.min_score,
        args.text_min_score,
        args.code_min_score
      ),
    ];

    const [originalResults, translatedResults] = await Promise.all(searchTasks);

    const allTextResults = [...originalResults.textResults, ...translatedResults.textResults];
    const allCodeResults = [...originalResults.codeResults, ...translatedResults.codeResults];

    let uniqueTextResults = Array.from(new Map(allTextResults.map(item => [item.id, item])).values())
      .sort((a, b) => b.score - a.score);

    let uniqueCodeResults = Array.from(new Map(allCodeResults.map(item => [item.id, item])).values())
      .sort((a, b) => b.score - a.score);

    // Reranking step
    if (args.rerank) {
      console.log("Reranking search results...");
      const reranker = new Reranker();

      // Rerank text results
      if (uniqueTextResults.length > 0) {
        const textDocuments = uniqueTextResults.map((result, index) => ({
          originalIndex: index,
          text: result.payload.codeChunk,
        }));
        const rerankedTextScores = await reranker.rerank(originalQuery, textDocuments);
        
        const rerankedTextResults = rerankedTextScores.map(reranked => {
          const originalResult = uniqueTextResults[reranked.originalIndex];
          originalResult.score = reranked.score; // Update score with relevance score
          return originalResult;
        });
        uniqueTextResults = rerankedTextResults.sort((a, b) => b.score - a.score);
      }

      // Rerank code results
      if (uniqueCodeResults.length > 0) {
        const codeDocuments = uniqueCodeResults.map((result, index) => ({
          originalIndex: index,
          text: result.payload.codeChunk,
        }));
        const rerankedCodeScores = await reranker.rerank(originalQuery, codeDocuments);

        const rerankedCodeResults = rerankedCodeScores.map(reranked => {
          const originalResult = uniqueCodeResults[reranked.originalIndex];
          originalResult.score = reranked.score; // Update score with relevance score
          return originalResult;
        });
        uniqueCodeResults = rerankedCodeResults.sort((a, b) => b.score - a.score);
      }
    }
    
    uniqueTextResults = uniqueTextResults.slice(0, args.max_results);
    uniqueCodeResults = uniqueCodeResults.slice(0, args.max_results);

    if (uniqueTextResults.length === 0 && uniqueCodeResults.length === 0) {
      return `🔍 No results found for query: "${originalQuery}" (and translated query: "${translatedQuery}")`;
    }
    
    let response = `🔍 Found ${uniqueTextResults.length + uniqueCodeResults.length} relevant snippets for: "${originalQuery}"\n\n`;

    if (uniqueTextResults.length > 0) {
      response += `## 📄 Text Files (${uniqueTextResults.length} results)\n\n`;
      uniqueTextResults.forEach((result, index) => {
        const payload = result.payload;
        if (payload) {
          const scoreLabel = args.rerank ? "Relevance" : "Score";
          response += `**${index + 1}. ${payload.filePath}** (lines ${payload.startLine}-${payload.endLine}) - ${scoreLabel}: ${result.score.toFixed(3)}\n`;
          response += "```\n" + payload.codeChunk + "\n```\n\n";
        }
      });
    }

    if (uniqueCodeResults.length > 0) {
      response += `## 💻 Code Files (${uniqueCodeResults.length} results)\n\n`;
      uniqueCodeResults.forEach((result, index) => {
        const payload = result.payload;
        if (payload) {
          const scoreLabel = args.rerank ? "Relevance" : "Score";
          response += `**${index + 1}. ${payload.filePath}** (lines ${payload.startLine}-${payload.endLine}) - ${scoreLabel}: ${result.score.toFixed(3)}\n`;
          response += "```\n" + payload.codeChunk + "\n```\n\n";
        }
      });
    }

    return response;

  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : String(error);
    return `❌ Search failed: ${errorMessage}`;
  }
}
